Skip to main content

🧠 Transfer Learning

You don't need millions of images to train a CNN!

🎓 The College Graduate

Models like ResNet spent weeks on supercomputers learning to identify 1,000 different objects. They are college graduates. If you want a model to classify X-Ray images, you don't start a new model from kindergarten. You hire the ResNet college graduate, chop off its final layer (the one that guesses the 1,000 objects), and add a new layer that guesses "Healthy" or "Sick".

Since it already knows how to detect edges and shapes, you can train it on just a few hundred X-Rays in a matter of minutes!

🐍 Python Implementation

import torch
import torch.nn as nn
import torchvision.models as models

# 1. Hire the graduate (pretrained=True loads the smart weights!)
model = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)

# 2. Freeze its brain (so we don't accidentally erase its memories)
for param in model.parameters():
param.requires_grad = False

# 3. Chop off the last layer (fc) and replace it with a new one
# We want it to only predict 2 things (Healthy vs Sick)
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, 2)

# Now only the new `fc` layer will learn during training!